#include "irblock.h"

#include <cstddef>
#include <functional>
#include <iostream>
#include <list>
#include <unordered_set>
#include <vector>

#include "arithmetic.h"
#include "ir.h"

using namespace Utils;

#define VAR BLUE
#define OP DEF
#define LABEL COLOR(5)

void IRBlock::AddIR(IR::IRBase* ir) {
    ir->parent = this;
    irList.push_back(ir);
}

void IRBlock::AddPhi(IR::Phi* phi) {
    phi->parent = this;
    phis.push_back(phi);
}

std::list<IR::IRBase*> IRBlock::GetAllIR() const {
    std::list<IR::IRBase*> ret;
    ret.insert(ret.end(), phis.begin(), phis.end());
    ret.insert(ret.end(), irList.begin(), irList.end());
    return ret;
}

std::vector<IRBlock*> IRBlock::GetNext() const {
    return GetTerminator()->GetNext();
}

std::vector<IRBlock*> IRBlock::GetPrev() const {
    std::vector<IRBlock*> ret;
    ret.reserve(prev.size());
    for (auto [u, _] : prev) ret.push_back(u);
    return ret;
}

IR::Terminator* IRBlock::GetTerminator() const {
    if (irList.empty()) return nullptr;
    auto ret = irList.back()->As<IR::Terminator>();
    return ret;
}

void IRBlock::Connect(IRBlock* nxt) {
    if (GetTerminator()) return;
    Assert(nxt->phis.empty(), "phi list should be empty");
    auto jmp = new IR::Jump();
    jmp->nxt = nxt;
    jmp->parent = this;
    irList.emplace_back(jmp);
    nxt->prev[this] = &jmp->nxt;
}

void IRBlock::Connect(ArithOp op, IR::Operand* lhs, IR::Operand* rhs,
                      IRBlock* t, IRBlock* f) {
    if (GetTerminator()) return;
    Assert(t->phis.empty() && f->phis.empty(), "phi list should be empty");
    auto br = new IR::Branch();
    br->op = op;
    br->lhs = lhs;
    br->rhs = rhs;
    br->t = t;
    br->f = f;
    br->parent = this;
    irList.emplace_back(br);
    t->prev[this] = &br->t;
    f->prev[this] = &br->f;
}

IRBlock* IRBlock::Remove(int reserve) {
    if (reserve == -1) reserve = GetNext()[0] == this;
    Assert(reserve < GetNext().size(), "");
    auto v = GetNext()[reserve];
    Assert(v->prev.count(this), "reverse edge not found");
    v->prev.erase(this);
    // u -> (this) -> v
    for (auto [u, ref] : prev) {
        Assert(*ref == this, "reverse edge not pointing to the right node");
        v->prev[u] = ref, *ref = v;
        for (auto phi : v->phis) {
            phi->srcs[u] = phi->srcs[this];
        }
    }
    for (auto phi : v->phis) {
        phi->srcs.erase(this);
    }
    return v;
}

IRBlock* IRBlock::RemoveBranch(int reserve) {
    if (reserve == -1) reserve = GetNext()[0] == this;
    auto w = GetNext()[!reserve];
    Assert(w->prev.count(this), "reverse edge not found");
    w->prev.erase(this);
    auto v = GetNext()[reserve];
    irList.pop_back();
    auto jmp = new IR::Jump();
    jmp->nxt = v;
    jmp->parent = this;
    irList.emplace_back(jmp);
    v->prev[this] = &jmp->nxt;
    return v;
}

void IRBlock::Print() const {
    std::cout << LABEL << name << ":\n";
    for (auto ir : phis) {
        std::cout << "    ", ir->Print(), std::cout << "\n";
    }
    for (auto ir : irList) {
        std::cout << "    ", ir->Print(), std::cout << "\n";
    }
}
